import os
import time

import pandas
import pandas as pd
import tushare as ts
from loguru import logger

from models.stock_model import StockNumber


def down_one(sn, today_date=None):
    try:
        sc = f'stocks/{sn.ts_code}.csv'
        if today_date and os.path.exists(sc):
            df2 = pandas.read_csv(sc)
            if str(df2['trade_date'][0]) != str(today_date):
                analysis_stock(sn)
            else:
                logger.info(f'{sn.name} 已经是最新数据,无需重新下载')
        else:
            analysis_stock(sn)
        return True
    except Exception as e:
        logger.error(e)
    return False


def get_all_stock(stocks, today_date=None):
    res = []
    df = pd.read_csv('all.csv')
    for row in df.index:
        if stocks is not None and df.loc[row]['ts_code'] not in stocks:
            continue
        sn = StockNumber(df.loc[row])
        if '退' in sn.name:
            continue
        # if 'ST' in sn.name:
        #     continue
        if not str(sn.ts_code).startswith('60') \
                and not str(sn.ts_code).startswith('30') \
                and not str(sn.ts_code).startswith('00'):
            continue
        # 下载中途停止，继续下载
        # if int(sn.ts_code.split('.')[0]) <= 603122:
        #     continue
        # count += 1
        try:
            sc = f'stocks/{sn.ts_code}.csv'
            if today_date and os.path.exists(sc):
                df2 = pandas.read_csv(sc)
                if str(df2['trade_date'][0]) != str(today_date):
                    analysis_stock(sn)
                else:
                    logger.info(f'{sn.name} 已经是最新数据,无需重新下载')
            else:
                analysis_stock(sn)
            res.append(sn.ts_code)
        except Exception as e:
            logger.error(e)
    return res


def analysis_stock(sn):
    """
    df = pro.daily(ts_code='000001.SZ', start_date='20180701', end_date='20180718')
    #多个股票
    df = pro.daily(ts_code='000001.SZ,600000.SH', start_date='20180701', end_date='20180718')
    """
    # 1分钟最低访问50次
    time.sleep(1.02)
    # 不复权
    # pro = ts.pro_api(token='bbc6f076aa3d2b063cb26376ad12ffd94b53864b42d2344b3aa2039f')
    # df2 = pro.daily(ts_code=sn.ts_code)
    # 前复权
    ts.set_token('bbc6f076aa3d2b063cb26376ad12ffd94b53864b42d2344b3aa2039f')
    df2 = ts.pro_bar(ts_code=sn.ts_code, adj='qfq')
    save_dir = f'../stocks'
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    df2.to_csv(f'stocks/{sn.ts_code}.csv')
    logger.info(f'save {sn.ts_code} {sn.name} 成功')


def run(stocks=None, today_date=None):
    """
    https://tushare.pro/document/2?doc_id=27
    接口：daily，可以通过数据工具调试和查看数据
    数据说明：交易日每天15点～16点之间入库。本接口是未复权行情，停牌期间不提供数据
    调取说明：120积分每分钟内最多调取500次，每次6000条数据，相当于单次提取23年历史
    描述：获取股票行情数据，或通过通用行情接口获取数据，包含了前后复权数据
    """
    logger.info(f'开始运行 download_all')
    res = get_all_stock(stocks, today_date)
    logger.success(f'运行完成 download_all')
    return res


if __name__ == '__main__':
    run()
